import java.awt.BorderLayout;
import java.awt.Container;
import java.awt.Dimension;
import java.awt.FlowLayout;
import java.awt.RenderingHints;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;

import javax.imageio.ImageIO;
import javax.swing.BorderFactory;
import javax.swing.ImageIcon;
import javax.swing.JFrame;
import javax.swing.JLabel;
import javax.swing.JPanel;

import java.util.ArrayList;
import java.util.concurrent.CyclicBarrier;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;


@SuppressWarnings("serial")
public class SegmentImg extends JPanel {
	public static int THUMB_MAX_WIDTH = 400;
	public static int THUMB_MAX_HEIGHT = 300;
	public static int SEGMENT_IMAGE_MAX_WIDTH = 400;
	public static int SEGMENT_MAGE_MAX_HEIGHT = 300;

	//default values for program arguments 
	static int numThreads=1; 
	static boolean parallelMode = false;
	static boolean displayOff = false;
	static boolean hold = false; 
	static String title = "Serial Segmentation";

	// Default image file names
	static String inputImageStr = "input.png";
	static String outputImageStr = "output";

	static String file_extension = "_sm."; //serial segmentation
	static String format = "png";

	private BufferedImage origThumb;

	private JLabel timerLabel;
	private JLabel timerLabelSeconds;
	private JLabel segmentLabel;
	private ImageIcon segmentedIcon;
	private Dimension segmentImageSize;

	private long time0;
	private long timeSub0;
	private long timeSub;

	private static long elapsedTime;

	public static CyclicBarrier barrier_finish, barrier_display;


	public SegmentImg(BufferedImage image) {

		if (!displayOff) {
			// create thumbnails
			segmentImageSize =  ImageUtils.determineSize(image.getWidth(), image.getHeight(), SEGMENT_IMAGE_MAX_WIDTH, SEGMENT_MAGE_MAX_HEIGHT);
			Dimension thumbSize = ImageUtils.determineSize(image.getWidth(), image.getHeight(), THUMB_MAX_WIDTH, THUMB_MAX_HEIGHT);
			this.origThumb = ImageUtils.getScaledInstance(image, thumbSize.width, thumbSize.height, RenderingHints.VALUE_INTERPOLATION_BILINEAR, true);

			this.createLayout();
		}
	}



	private void createLayout() {
		this.setLayout(new BorderLayout());

		//Display a thumbnail of the original image
		Container topContainer = new JPanel();
		topContainer.setLayout(new FlowLayout(FlowLayout.CENTER));
		Container btmContainer = new JPanel();
		btmContainer.setLayout(new FlowLayout(FlowLayout.CENTER));
		ImageIcon fromIcon = new ImageIcon(origThumb);
		JLabel fromLabel = new JLabel("");
		fromLabel.setIcon(fromIcon);
		fromLabel.setBorder(BorderFactory.createTitledBorder("Original image"));
		topContainer.add(fromLabel);

		this.add(topContainer, BorderLayout.WEST);

		//Display the image being changed

		segmentLabel = new JLabel();
		segmentLabel.setBorder(BorderFactory.createTitledBorder("Segmented image"));
		btmContainer.add(segmentLabel);
		this.add(btmContainer, BorderLayout.EAST);

		Container timerContainer = new JPanel();
		timerLabel = new JLabel("Time taken: ");
		timerContainer.add(timerLabel);
		timerLabelSeconds = new JLabel("0 ms");
		timerContainer.add(timerLabelSeconds);
		this.add(timerContainer, BorderLayout.SOUTH);
	}

	public void startTimer() {
		time0 = System.currentTimeMillis();
		timeSub = 0;
		timeSub0 = 0;
	}

	public void stopTimer() {
		elapsedTime = System.currentTimeMillis() - time0 - timeSub;
	}

	public void freezeTimer() {
		timeSub0 = System.currentTimeMillis();
	}

	public void unFreezeTimer() {
		timeSub += System.currentTimeMillis() - timeSub0;
	}

	public void updateTimer(int phase) {
		if (!displayOff) {
			long t = System.currentTimeMillis() - time0;
			t -= timeSub;
			if(phase !=-1)
				timerLabelSeconds.setText(Long.toString(t) + " ms" + "  (Phase: " + phase + ")" );
			else 
				timerLabelSeconds.setText(Long.toString(t) + " ms");	
		}
	}

	public void updateDoneLabel() {
		if (!displayOff)
			timerLabel.setText("Image segmentation finished in: ");
	}

	public void updateSegmentedImage(BufferedImage image, boolean hiQuality) {
		if (!displayOff) {
			BufferedImage displayImage = ImageUtils.getScaledInstance(image, segmentImageSize.width, segmentImageSize.height, RenderingHints.VALUE_INTERPOLATION_BILINEAR, hiQuality);
			segmentedIcon = new ImageIcon(displayImage);
			segmentLabel.setIcon(segmentedIcon);
		}
	}




	static void Usage() {
		System.err.println("Usage: SegmentImg \n" +
				"[-f <format> ] use format as output format. i.e. png. (default png)\n" + 
				"[-p <n> ] switch to parallel mode with n number of threads\n" + 
				"[-s <n> ] threshold value (default 4)\n" +
				"[-R <n> ] use <n> as seed to the random number generator which in turn chooses color on labels.\n" +
				"[-t ] turn off display\n" + 
				"[-h ] hold the display\n" + 
		"[-i inputimage ]\n");

		System.exit(0);
	}

	static void parseInputArguments(String args[])
	{

		int i = 0;
		int threshold = 4;

		// Parse input parameters
		while (i < args.length && args[i].startsWith("-")) 
		{

			String arg = args[i++];

			if (arg.equals("-i")) {
				// input file
				inputImageStr = args[i++];                 
				outputImageStr = inputImageStr.substring(0, inputImageStr.lastIndexOf('.')) + "_out";

			}
			// Parallel mode
			else if (arg.equals("-p")) {
				parallelMode = true;
				try{ numThreads = ( Integer.parseInt(args[i++]));
				} catch (NumberFormatException Exc){
					System.err.println("-p needs an int");
					Usage();
				}
				Segmentation.numThreads = numThreads;
				title = "Parallel Segmentation";
				// The output file will have the extension "_pm" added
				file_extension = "_pm."; // parallel segmentation

			}
			else if (arg.equals("-s")) {
				try{ threshold = ( Integer.parseInt(args[i++]));
				} catch (NumberFormatException Exc){
					System.err.println("-s needs an int");
					Usage();
				}
				Segmentation.threshold = threshold;
			}
			else if (arg.equals("-h")) {
				hold = true;
			}
			else if (arg.equals("-f")) {
				String outputformat = args[i++];
				String[] supportedformatnames = ImageIO.getWriterFormatNames();

				if (!Arrays.asList(supportedformatnames).contains(outputformat)) {
					System.err.println("format " + outputformat + " not supported. supported formats are:");
					System.err.println(Arrays.asList(supportedformatnames));
					Usage();
				}
				format = outputformat;
			}
			else if (arg.equals("-R")) {
				int randomseed = -1;
				try{
					randomseed = ( Integer.parseInt(args[i++]));
				} catch (NumberFormatException Exc){
					System.err.println("-R needs an int");
					Usage();
				}
				Segmentation.randomSeed = randomseed;
			}

			else if (arg.equals("-t")) {
				displayOff = true;
				System.out.println("Display turned off!");
				SegmentImg.displayOff = true;
			}
			else{
				System.err.println("Unknown option " + arg);
				Usage();
			}
		}
		if (i < args.length) {
			System.err.println("Unknown option " + args[i]);
			Usage();
		}
	}

	public static void gatherStatistics(int [] labels)
	{
		HashMap<Integer,Integer> clusters= new HashMap<Integer,Integer>();

		int min_size = labels.length ; // there is only one cluster
		int max_size = 1;   // every pixel is a cluster

		int total_size = 0 ; 

		for(int k=0 ; k < labels.length ; k++)
		{
			if (labels[k] == 0) {
				continue; // this is a background label
			}
			if(!clusters.containsKey(labels[k]))
				clusters.put(labels[k], 1);  // for the first time, we encounter such a label
			else
				clusters.put(labels[k], clusters.get(labels[k]) + 1 );	

		}

		Collection<Integer> c = clusters.keySet();

		Iterator<Integer> itr = c.iterator();


		while(itr.hasNext()){

			int key = itr.next();

			int count = clusters.get(key);

			min_size = Math.min(min_size, count);
			max_size = Math.max(max_size, count);

			total_size += count; 
		}

		System.out.println("Number of clusters found: "  +   clusters.size() ) ;
		System.out.println("Min cluster size        : "  +  min_size) ;
		System.out.println("Max cluster size        : "  +  max_size) ;
		System.out.println("Average cluster size    : "  +  ((double)
				total_size / (double) clusters.size())) ;

	}


	public static void main(String[] args) {

		JFrame frame;

		parseInputArguments(args);

		File imgFile = new File(inputImageStr);
		assert(imgFile.exists() && imgFile.isFile());

		//Read images
		BufferedImage image = null;
		try {
			image = ImageIO.read(imgFile);
		} catch(IOException e) {
			e.printStackTrace();
		}

		// Initiate gui 
		final SegmentImg sa = new SegmentImg(image);
		if (!displayOff) {
			frame = new JFrame(title);
			frame.setContentPane(sa);
			frame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
			frame.pack();
			frame.setVisible(true);
			frame.setSize(800, 400);
		}


		/* The parallel section */
		if (parallelMode) {

			
			// create representation of image; have to split up the image in accordance to the number of threads:

			// to account for the number of threads. Each thread gets a strip of
			// the image to analyze.
			int height = image.getHeight();
			int width = image.getWidth();
			System.out.printf("Load image %d x %d\n", width, height);
			int work_remainder = height % numThreads;
			int work_divisible = height - work_remainder;
			
			final ArrayList<ParallelSegmentation> threads = new ArrayList<ParallelSegmentation>(numThreads);
			//barrier to getting all the threads done with running, then show the output
			
			final int labels[] = new int[width*height];
			barrier_finish = new CyclicBarrier(numThreads, new Runnable() {
				public void run() { 
					System.out.println("in barrier");

					sa.stopTimer();

					gatherStatistics(labels);

					sa.updateTimer(-1);

					BufferedImage segmentedImage = threads.get(0).getBufferedImage();
					sa.updateSegmentedImage(segmentedImage, true);

					// Write output

					File output = new File( outputImageStr + file_extension  + format);
					try {
						ImageIO.write(segmentedImage, format, output);
					} catch (IOException e) {
						e.printStackTrace();
					}

					sa.updateSegmentedImage(segmentedImage, true);
					System.out.println("Elapsed Time (ms): " +  Long.toString(elapsedTime));

					if(!displayOff){
						try{
							if(!hold){
								// last 4 sec before self destruction
								Thread.sleep(4000); 
								//System.exit(0);
							}
						}
						catch(Exception e)
						{
							System.out.println(e);
						}
					}
				}
			});

			sa.startTimer();
			ParallelSegmentation.time_start = System.currentTimeMillis();
			int imagePixels[] = image.getRaster().getPixels(0, 0, width, height, (int[])null);




			// Space between two pixels in the imagePixel array.
			int pixelWidth = image.getSampleModel().getNumDataElements();

			// Apply initial labels (each label is it's own index in the label array + 1)
			

			for (int i = 0; i < width*height; i++) { 
				labels[i] = 0; 
			}

			// Label each pixel as a separate label.
			for (int i = 0; i < height; i++) {
				for (int j = 0; j < width; j++) {

					int idx = (i*width + j);

					int idx3 = idx*pixelWidth;
					// Comment this line if you want to label background pixels
					if (imagePixels[idx3] == 0) 
						continue;
					labels[idx] = idx+1;

				}
			}	
			
			BufferedImage segImage = new BufferedImage(width, height, BufferedImage.TYPE_3BYTE_BGR);
			for (int i = 0; i < numThreads; i++)
			{
				int start = work_divisible/numThreads * i;
				int finish = (work_divisible)/numThreads * (i+1);

				// first thread gets the extra work
				if (i != 0)
					start += work_remainder;
				
				finish += work_remainder;
				
				ParallelSegmentation parallelsegmentation = 
					new ParallelSegmentation(image, segImage, sa, labels, start, finish);


				parallelsegmentation.setId(i);
				parallelsegmentation.start();

				threads.add(parallelsegmentation);



			}


		} else {
			Segmentation segmentation = new Segmentation();
			segmentation.SegmentationInit(image, sa);
			sa.startTimer();

			int []labels = segmentation.doSegmentation();

			sa.stopTimer();

			gatherStatistics(labels);

			sa.updateTimer(-1);

			BufferedImage segmentedImage = segmentation.getSegmentedImage();
			sa.updateSegmentedImage(segmentedImage, true);

			// Write output

			File output = new File( outputImageStr + file_extension  + format);
			try {
				ImageIO.write(segmentedImage, format, output);
			} catch (IOException e) {
				e.printStackTrace();
			}

			System.out.println("Elapsed Time (ms): " +  Long.toString(elapsedTime));

			if(!displayOff){
				try{
					if(!hold){
						// last 4 sec before self destruction
						Thread.sleep(4000); 
						//System.exit(0);
					}
				}
				catch(Exception e)
				{
					System.out.println(e);
				}
			}

		}

	}
}

